import sys
import os
import subprocess

def print_usage():
    print("""
Usage: python main.py <operation_number>

Available operations:
1: Clean previous content (clean.py)
2: Merge pictures and perform OCR (merge_sortpic.py + llm_ocr.py)
3: Extract dialogues (extract_dialogues.py)
4: Extract manga panels (manga_panel_extractor.py)
5: Organize panels (organize_panels.py)
6: Generate video (video_generator.py)

Example: python main.py 2
    """)

def activate_conda_env_py312():
    # Get the conda activation command
    if sys.platform == "win32":
        conda_activate = "conda activate py312"
    else:
        conda_activate = "source $(conda info --base)/etc/profile.d/conda.sh && conda activate py312"
    
    return conda_activate

def activate_conda_env_base():
    # Get the conda activation command
    if sys.platform == "win32":
        conda_activate = "conda activate base"
    else:
        conda_activate = "source $(conda info --base)/etc/profile.d/conda.sh && conda activate base"
    
    return conda_activate

def activate_conda_env_py39():
    # Get the conda activation command
    if sys.platform == "win32":
        conda_activate = "conda activate py39"
    else:
        conda_activate = "source $(conda info --base)/etc/profile.d/conda.sh && conda activate py39"
    
    return conda_activate

def run_script(script_name, operation_number):
    try:
        # Select conda environment based on operation number
        if operation_number in [1, 2, 4, 5]:
            conda_activate = activate_conda_env_py312()
        elif operation_number == 6:
            conda_activate = activate_conda_env_base()
        elif operation_number == 7:
            conda_activate = activate_conda_env_py39()
        else:
            conda_activate = activate_conda_env_py312()  # default to py312

        # Combine conda activation with the script execution
        full_command = f"{conda_activate} && python {script_name}"
        subprocess.run(full_command, shell=True, check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error running {script_name}: {e}")
        sys.exit(1)
    except FileNotFoundError:
        print(f"Script {script_name} not found!")
        sys.exit(1)

def main():
    if len(sys.argv) != 2:
        print_usage()
        sys.exit(1)

    try:
        operation = int(sys.argv[1])
    except ValueError:
        print("Error: Please provide a valid number")
        print_usage()
        sys.exit(1)

    # Dictionary mapping operation numbers to functions
    operations = {
        1: lambda: run_script('clean.py', 1),
        2: lambda: (run_script('merge_sorted_images.py', 2), run_script('llm_ocr.py', 2), run_script('extract_dialogues.py', 2)),#台词生成，需要去改一下dify中的提示词
        # 3: lambda: run_script('extract_dialogues.py', 3),
        4: lambda: run_script('manga_panel_extractor.py', 4),#分镜处理
        5: lambda: run_script('organize_panels.py', 5),
        6: lambda: run_script('test_ocr.py', 6), #把台词放入对应文件夹中
        7: lambda: run_script('video_generator.py', 7),
        8: lambda: run_script('manga_panel_extractor_v2.py', 4),#分镜处理
        9: lambda: run_script('manga_panel_extractor_black_line.py', 4),#黑线分镜处理
    }

    if operation not in operations:
        print(f"Error: Operation {operation} is not valid")
        print_usage()
        sys.exit(1)

    # Execute the selected operation
    operations[operation]()

if __name__ == "__main__":
    main() 